import numpy as np

JOINT_MAX = np.array([150, 180, 0, 100, 70, 120])
JOINT_MIN = np.array([-150, 0, -170, -100, -70, -120])

# Estimated from the data
END_MAX = np.array([752, 704, 877, 180,  90, 180])
END_MIN = np.array([-642, -741,  -18, -180,  -26, -180])


def normalize_joint(joints):
    return (joints - JOINT_MIN) * 2 / (JOINT_MAX - JOINT_MIN) - 1


def denormalize_joint(normed_joints):
    # r_xyz.shape = (B, 6)
    return (normed_joints + 1) / 2 * (JOINT_MAX - JOINT_MIN) + JOINT_MIN


def normalize_end(pos):
    dim = pos.shape[-1]
    return (pos - END_MIN[:dim]) * 2 / (END_MAX[:dim] - END_MIN[:dim]) - 1


def denormalize_end(normed_pos):
    dim = normed_pos.shape[-1]
    return (normed_pos + 1) / 2 * (END_MAX[:dim] - END_MIN[:dim]) + END_MIN[:dim]


def inverse_fourier_joint(f_joints):
    # f_joints.shape = (B, 12)
    cos_j = f_joints[:, ::2]
    sin_j = f_joints[:, 1::2]
    return np.sign(sin_j) * np.arccos(cos_j) / np.pi * 180
